
import numpy as np
from panda3d.core import OmniBoundingVolume
from panda3d.core import DirectionalLight, AmbientLight, Material
from panda3d.core import ShaderAttrib, Shader, ShaderBuffer
from panda3d.core import Vec3, Vec4, Point3, LPoint3
from panda3d.core import GeomNode, Geom
from panda3d.core import NodePath
from panda3d.core import BitMask32  
from panda3d.bullet import BulletBoxShape, BulletRigidBodyNode
from panda3d.core import ColorBlendAttrib, TransparencyAttrib

from itertools import product

def intv(v):
    x, y, z = v
    # panda3d的低精度向量运算有误差，int()有时会相差1
    return round(x), round(y), round(z)


class BlockManager:
    def __init__(self, world, bulletworld):
        self.chunk_data = {}
        self.world = world
        self.bulletworld = bulletworld
        self.box = loader.loadModel('models/box2.egg').node()
        self.mysh = Shader.load(Shader.SL_GLSL, "shaders/test.vert", "shaders/test.frag")
        self.mysh2 = Shader.load(Shader.SL_GLSL, "shaders/test.vert", "shaders/glass.frag")
        
        self.physics_interest = {}      #{int: (x, y, z)}
        self.physics_dirty = True
        self.physics_interest_off = {} # {i: p}
        self.physics_data = {}
        
        # ~ shape = BulletBoxShape(Vec3(0.5, 0.5, 0.5))
        # ~ shape = BulletBoxShape(Vec3(0.48, 0.48, 0.48))
        shape = BulletBoxShape(Vec3(0.48, 0.48, 0.42))
        n = BulletRigidBodyNode('Box')                            # 这个碰撞体的运动会影响父级的位置啊
        np = NodePath(n)
        np.setCollideMask(BitMask32(0b00011001))     # block mask
        n.addShape(shape)
        self.help_box = np
        self.help_boxn = n
        
    
    @staticmethod
    def get_near(p):
        if p is None:
            return set()
        else:
            x, y, z = p
            return set(product(range(x-2, x+2), range(y-2, y+2), range(z-2, z+2)))
        
    
    @staticmethod
    def divmod(p):
        x, y, z = p
        bx, x = divmod(x, 32)
        by, y = divmod(y, 32)
        bz, z = divmod(z, 32)
        return (bx, by, bz), (x, y, z)
        
    def setblock(self, p, type, flush):
        self.physics_dirty = True
        
        bp, p = self.divmod(p)
        # ~ print("Big set", bp, p)
        if bp in self.chunk_data:
            self.chunk_data[bp].setblock(p, type, flush)
        else:
            cd = BigChunk(self, bp)
            self.chunk_data[bp] = cd
            cd.setblock(p, type, flush)
    
    def delblock(self, p, flush):
        if p in self.physics_data:
            np = self.physics_data[p]
            self.bulletworld.removeRigidBody(np.node())
            np.remove_node()
            del self.physics_data[p]
        
        
        bp, p = self.divmod(p)
        if bp in self.chunk_data:
            self.chunk_data[bp].delblock(p, flush)
        else:
            print("delblock error: not exist")
            pass
            
        

            
    def getblock(self, p):
        bp, p = self.divmod(p)
        if bp in self.chunk_data:
            return self.chunk_data[bp].getblock(p)
        else:
            return None
    
    def flush(self):
        for v in self.chunk_data.values():
            v.prepare()
        
    
    def physics_interest_new(self, p):
        i = 1
        while i in self.physics_interest:
            i += 1
            
        p = intv(p)
        self.physics_interest[i] = p
        self.physics_interest_off[i]= None
        self.physics_dirty = True
        print("new physics interest")
        return i
        
    def physics_interest_del(self, i):
        if i in self.physics_interest:
            del self.physics_interest[i]
        else:
            print("physics_interest_del error, ", i, "not exsits")
            
    def physics_interest_update(self, i, p):
        if i in self.physics_interest:
            p = intv(p)
            p_old = self.physics_interest[i]
            if p != p_old:
                print("move check", p, p_old)
                self.physics_interest[i] = p
                if i not in self.physics_interest_off:
                    self.physics_interest_off[i] = p_old
                self.physics_dirty = True
        else:
            print("physics_interest_update error, ", i, "not exsits")
        
    
    def physics_render(self):
        if self.physics_dirty == False:
            # ~ print("no render")
            return
        else:
            # ~ print("real render")
            all_new = set()
            all_old = set()
            all_static = set()
            for new in self.physics_interest.values():
                this_new = self.get_near(new)
                all_new.update(this_new)
                
            for old in self.physics_interest_off.values():
                this_old = self.get_near(old)
                all_old.update(this_old)
            
            
            pos = all_new
            neg = all_old - all_new
            
            for p in neg:
                if p in self.physics_data:
                    np = self.physics_data[p]
                    self.bulletworld.removeRigidBody(np.node())
                    np.remove_node()
                    del self.physics_data[p]
                else:
                    pass
                    # ~ print("physics render, remove block 未检测到", p)
            for p in pos:
                if p in self.physics_data:
                    pass
                    # ~ print("physics render, add block 重复")
                else:
                    if self.getblock(p) is None:
                        continue 
                    shape = BulletBoxShape(Vec3(0.5, 0.5, 0.5))
                    # ~ shape = BulletBoxShape(Vec3(0.499, 0.499, 0.499))
                    n = BulletRigidBodyNode('Box')                            # 这个碰撞体的运动会影响父级的位置啊
                    n.set_friction(0)
                    # ~ n.set_restitution(0.09)
                    np = NodePath(n)
                    n.set_mass(0.0)
                    np.setCollideMask(BitMask32(0b00011001))     # block mask
                    n.addShape(shape)
                    n. set_python_tag("block", None)
                    self.bulletworld.attach_rigid_body(n)
                    x, y, z = p
                    v = Vec3(x + 0.5, y + 0.5, z + 0.5)
                    np.set_pos(v)
                    self.physics_data[p] = np
                    
            
            self.physics_dirty = False
            self.physics_interest_off.clear()
            # ~ print(self.physics_dirty, "dirty")
        

class BigChunk:
    def __init__(self, mgr, bp):
        self.mgr = mgr
        self.bp = bp
        self.chu_data = {}
        self.dirty = True
        
        
        
        
        if "scene data":
            myMaterial = Material()
            myMaterial.setShininess(5.0)
            myMaterial.setAmbient((1, 1, 1, 1)) #Make this material blue
            
            K = 0.3
            alight = AmbientLight('ambientLight')
            alight.set_color(Vec4(K, K, K, 1))
            
            
            K = 0.8
            dlight1 = DirectionalLight('directionalLight')
            dlight1.set_direction(Vec3(1, 1.6, -0.65))
            dlight1.set_color(Vec4(K, K, K, 1))
            
            dlight2 = DirectionalLight('directionalLight2')
            dlight2.set_direction(Vec3(-1, -1.6, -0.65))
            dlight2.set_color(Vec4(K, K, K, 1))
        
        
        
        # ~ self.box = mgr.box.__copy__()
        self.box = NodePath(mgr.box.__copy__())
        # ~ self.box = loader.loadModel('models/box2.egg')
        self.box.set_pos(Vec3(0, 0, 0))
        self.box.set_shader_auto(False)
        self.box.set_shader(mgr.mysh)
        
        
        self.sb = ShaderBuffer("te", np.zeros((0, 3), dtype="f4"), Geom.UHStatic)
        self.box.set_shader_input("te", self.sb)
        self.box.set_instance_count(0)
        # ~ self.box.hide()
        
        self.box.set_shader_input("attr_light0_color",  dlight1.get_color())
        self.box.set_shader_input("attr_light0_vec", dlight1.get_direction())
        self.box.set_shader_input("attr_light1_color", dlight2.get_color())
        self.box.set_shader_input("attr_light1_vec", dlight2.get_direction())
        self.box.set_shader_input("attr_material_ambient", myMaterial.get_ambient())
        self.box.set_shader_input("attr_ambient", alight.get_color())
        self.box.set_shader_input("attr_colorscale", 1.0)
        # ~ self.box.set_shader_input("te", self.sb)
        # ~ self.box.set_instance_count(0)
        
        self.box.node().setBounds(OmniBoundingVolume())
        self.box.node().setFinal(True)
        
        ###
        self.box.reparent_to(self.mgr.world)
        
        
        if "glass" :
            self.box2 = NodePath(mgr.box.__copy__())
            # ~ self.box2 = loader.loadModel('models/box2.egg')
            
            self.box2.set_pos(Vec3(0, 0, 0))
            self.box2.set_shader_auto(False)
            self.box2.set_shader(mgr.mysh2)
        
            self.sb2 = ShaderBuffer("te", np.zeros((0, 3), dtype="f4"), Geom.UHStatic)
            self.box2.set_shader_input("te", self.sb2)
            self.box2.set_instance_count(0)
        
            self.box2.set_shader_input("attr_light0_color",  dlight1.get_color())
            self.box2.set_shader_input("attr_light0_vec", dlight1.get_direction())
            self.box2.set_shader_input("attr_light1_color", dlight2.get_color())
            self.box2.set_shader_input("attr_light1_vec", dlight2.get_direction())
            self.box2.set_shader_input("attr_material_ambient", myMaterial.get_ambient())
            self.box2.set_shader_input("attr_ambient", alight.get_color())
            self.box2.set_shader_input("attr_colorscale", 1.0)
            
            self.box2.reparent_to(self.mgr.world)
            
            self.box2.node().setBounds(OmniBoundingVolume())
            self.box2.node().setFinal(True)
            
            self.box2.set_transparency(TransparencyAttrib.MAlpha)
                
        self.setblock((0, 0, -1000), 1, flush=True)
        self.setblock((0, 0, -1001), 2, flush=True)
        
        
    
    @staticmethod
    def divmod(p):
        x, y, z = p
        bx, x = divmod(x, 4)
        by, y = divmod(y, 4)
        bz, z = divmod(z, 4)
        return (bx, by, bz), (x, y, z)
    

        
    
    def setblock(self, p, type, flush):
        
        bp, p = self.divmod(p)
        # ~ print("small set", bp, p)
        if bp in self.chu_data:
            self.chu_data[bp].setblock(p, type)
        else:
            cd = SmallChunk(self, bp)
            self.chu_data[bp] = cd
            cd.setblock(p, type)
        
        self.dirty = True
        if flush:
            self.prepare()
    
    def delblock(self, p, flush):
        bp, p = self.divmod(p)
        if bp in self.chu_data:
            r = self.chu_data[bp].delblock(p)
            if r:
                self.dirty = True
        else:
            pass
            
        if flush:
            self.prepare()
            
    def getblock(self, p):
        bp, p = self.divmod(p)
        if bp in self.chu_data:
            return self.chu_data[bp].getblock(p)
        else:
            return None
    
        # 根据chu_data来计算渲染的方块
    def prepare(self):
        if self.dirty is False:
            return
            
        if self.sb is not None:
            self.sb.release_all()
        if self.sb2 is not None:
            self.sb2.release_all()
        
        acc1 = []
        acc2 = []
        for k, v in self.chu_data.items():
            bp = np.array([k[0], k[1], k[2], 0]) * 4
            l1, l2 = v.collect()
            acc1.append(l1 + bp)
            acc2.append(l2 + bp)
        l1 = np.concatenate(acc1, dtype="f4")
        bp = self.bp
        l1 += np.array([bp[0], bp[1], bp[2], 0]) * 32
        l1+= np.array([0.5, 0.5, 0.5, 0])
        
        l2 = np.concatenate(acc2, dtype="f4")
        bp = self.bp
        l2 += np.array([bp[0], bp[1], bp[2], 0]) * 32
        l2 += np.array([0.5, 0.5, 0.5, 0])
        
        
        self.sb = ShaderBuffer("te", l1, Geom.UHStatic)
        self.box.set_shader_input("te", self.sb)
        self.box.set_instance_count(l1.shape[0])
        
        self.sb2 = ShaderBuffer("te", l2, Geom.UHStatic)
        self.box2.set_shader_input("te", self.sb2)
        self.box2.set_instance_count(l2.shape[0])
        
        
        self.dirty = False
    

class SmallChunk:
    def __init__(self, big, bp):
        self.big = big
        self.bp = bp
        self.block_data = {}
        self.dirty = True
        self.cache = None
        
        
    @staticmethod
    def divmod(p):
        x, y, z = p
        bx, x = divmod(x, 4)
        by, y = divmod(y, 4)
        bz, z = divmod(z, 4)
        return (bx, by, bz), (x, y, z)
    
    
    def setblock(self, p, type):
        if p in self.block_data:
            self.block_data[p] = type
        else:
            self.block_data[p] = type
            
            
                
        
            
        self.dirty = True
        
    def delblock(self, p):
        if p in self.block_data:
            del self.block_data[p]
            self.dirty = True
            return True
        else:
            return False
            
    def getblock(self, p):
        if p in self.block_data:
            assert self.block_data[p] is not None
            return self.block_data[p]
        else:
            return None
            
        
        
    def collect(self):      # -> (l1, l2)
        if self.dirty is False:
            return self.cache
        
        l1 = []
        l2 = []
        for k, v in self.block_data.items():
            if v == 1:
                l1.append(k)
            elif v == 2:
                l2.append(k)
        
        
        n1 = len(l1)
        if n1 == 0:
            l1 = np.zeros((0, 3), dtype="f4")
        else:
            l1 = np.array(l1, dtype="f4")
        n1 = l1.shape[0]
        l1 = np.concatenate([l1, np.zeros((n1, 1))], axis=1, dtype="f4")
        n2 = len(l2)
        if n2 == 0:
            l2= np.zeros((0, 3), dtype="f4")
        else:
            l2 = np.array(l2, dtype="f4")
        l2 = np.concatenate([l2, np.zeros((n2, 1))], axis=1, dtype="f4")
        
        self.cache = (l1, l2)
        self.dirty = False
        
        return (l1, l2)
    
    
        
